import numpy as np
import matplotlib.pyplot as plt

# ---------------------- 1. 原始数据整理 ----------------------
# 因变量：反电动势（施加的电压V）
v = np.array([1.0, 1.5, 2.0, 2.5, 3.0])
# 自变量：平均转速ω（rad/s）
omega = np.array([346.2, 508.5, 642.1667, 709.6667, 802.0])

# ---------------------- 2. 三次多项式拟合（强制常数项为0） ----------------------
# 构造拟合矩阵：X = [ω³, ω², ω]（无常数项）
X = np.vstack([omega**3, omega**2, omega]).T
# 最小二乘法拟合：y = a*ω³ + b*ω² + c*ω
coeffs = np.linalg.lstsq(X, v, rcond=None)[0]
a, b, c = coeffs  # 提取三次项、二次项、一次项系数

# 输出拟合系数
print("三次多项式拟合结果（E = a*ω³ + b*ω² + c*ω）：")
print(f"三次项系数 a = {a:.9f}")
print(f"二次项系数 b = {b:.7f}")
print(f"一次项系数 c = {c:.6f}")

# ---------------------- 3. 生成拟合曲线，验证效果 ----------------------
# 生成更多转速点，绘制拟合曲线
omega_fit = np.linspace(300, 850, 100)  # 转速范围：300~850 rad/s
v_fit = a * omega_fit**3 + b * omega_fit**2 + c * omega_fit

# 绘图
plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文显示
plt.figure(figsize=(10, 6))
# 原始数据点
plt.scatter(omega, v, color='red', s=80, label='实测数据（电压-平均转速）')
# 拟合曲线
plt.plot(omega_fit, v_fit, color='blue', linewidth=2, label='三次多项式拟合曲线')
plt.xlabel('电机转速 ω (rad/s)')
plt.ylabel('反电动势 E (V)')
plt.title('电机反电动势-转速 三次多项式拟合')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

# ---------------------- 4. 输出拟合函数（直接复制到C代码中） ----------------------
print("\n对应C语言反电动势函数：")
print(f"float Motor_CalcRevVoltXXX(float speed) {{")
print(f"    return {a:.9f}f * speed * speed * speed + {b:.7f}f * speed * speed + {c:.6f}f * speed;")
print(f"}}")